package com.tianji.aigc.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.extension.conditions.query.LambdaQueryChainWrapper;
import com.tianji.aigc.config.SessionProperties;
import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.memory.RedisChatMemory;
import com.tianji.aigc.result.ToolResultHolder;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BizIllegalException;
import com.tianji.common.utils.UserContext;
import org.aspectj.weaver.ast.Var;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

@Service
public class ChatServiceImpl implements ChatService {

    @Autowired
    ChatClient chatClient;
    
    
    @Autowired
    private IChatSessionService chatSessionService;

    @Autowired
    private RedisChatMemory redisChatMemory;
    
    
    private static Map<String,Boolean> SESSION_MAP=new ConcurrentHashMap<String,Boolean>();

    @Autowired
    private VectorStore vectorStore;

    @Autowired
    SessionProperties sessionProperties;
    

    /**
     * 聊天接口
     * @param chatDTO
     * @return
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO chatDTO) {

        ChatSession chatSession = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, chatDTO.getSessionId()).one();
        if(chatSession==null){
            throw new BizIllegalException("会话不存在");
        }
         //更新第一次对话的会话标题
        chatSessionService.updateSessionTitle(chatDTO.getQuestion(),chatDTO.getSessionId(),UserContext.getUser());
        String converationId = UserContext.getUser()+":"+chatDTO.getSessionId();
        StringBuffer content=new StringBuffer();

        //生成唯一Id,大map的Id
        String requestId= IdUtil.fastSimpleUUID();

        return chatClient.prompt()
                .user(chatDTO.getQuestion())
                .toolContext(MapUtil.<String,Object>builder().put("request_id",requestId)
                        .put("user_id",UserContext.getUser())
                        .build())
                .advisors(x->x
                        .advisors(new QuestionAnswerAdvisor(vectorStore,SearchRequest.builder().query("").topK(999).build()))
                        .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,converationId))
                .stream()
                //.content()
                .chatResponse()
                .doFirst(()-> SESSION_MAP.put(chatDTO.getSessionId(),true))//标记开始输出的第一步，设置为true，表示可以继续输出
                .doOnComplete(()-> SESSION_MAP.remove(chatDTO.getSessionId()))//标记输出完成，清除标记
                .doOnError(x-> SESSION_MAP.remove(chatDTO.getSessionId()))//输出异常，清除标记
                .doOnCancel(()->this.saveRedis(converationId, content.toString()))//发现取消输出时（输出停止时）将之前收集到的全部消息保存到Redis
                .takeWhile(x->SESSION_MAP.getOrDefault(chatDTO.getSessionId(),false))//根据布尔值决定是否继续输出，true继续，false停止
                .map(x-> {
                    //判断会话结束的原因
                    String finishReason = x.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals("STOP",finishReason)){
                        //获取消息id
                        String messageId = x.getMetadata().getId();
                        //将消息Id与requestId关联
                        ToolResultHolder.put(messageId,"request_id",requestId);
                    }

                    String text = x.getResult().getOutput().getText();
                    content.append(text);
                    return ChatEventVO.builder().eventData(text).eventType(ChatEventTypeEnum.DATA.getValue()).build();

                })
                .concatWith(Flux.defer(()->{
                    Map<String, Object> resultMap = ToolResultHolder.get(requestId);
                    if (CollUtil.isNotEmpty(resultMap)) {
                        ChatEventVO chatEventVO = ChatEventVO.builder().eventData(resultMap).eventType(ChatEventTypeEnum.PARAM.getValue()).build();

                        ToolResultHolder.remove(requestId);
                        return Flux.just(chatEventVO,ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build());
                    }

                    return Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build());
                }));

    }

    private void saveRedis(String converationId, String content) {
        redisChatMemory.add(converationId,new AssistantMessage(content));
    }
    /*Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build())*/

    /**
     * 停止生成
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId,false);
    }
    /**
     * 文本聊天功能
     */
    @Override
    public String TextChat(String question) {

        return chatClient.prompt()
               .system(x->x.text(sessionProperties.getText()))
               .user(question)
               .call()
               .content();
    }
}
